﻿using FaissMask;
using System;
using System.Collections.Generic;
using System.Diagnostics.PerformanceData;
using System.Linq;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;

namespace FaissOperation
{
    /// <summary>
    /// 矢量知识库
    /// </summary>
    public class VKB
    {
        private Dictionary<string, Cluster> _clusters = new Dictionary<string, Cluster>();

        private static VKB __instance = new VKB();

        private VKB() { }

        /// <summary>
        /// 获取VKB的单例实例。
        /// </summary>
        /// <returns>VKB的单例实例。</returns>
        public static VKB getInstance()
        {
            return __instance;
        }

        /// <summary>
        /// 创建或获取指定键的Cluster实例。
        /// </summary>
        /// <param name="key">Cluster的键。</param>
        /// <returns>对应的Cluster实例。</returns>
        public Cluster CreateCluster(string key)
        {
            if (_clusters.ContainsKey(key))
            {
                return _clusters[key];
            }
            else
            {
                Cluster cluster = new Cluster();
                _clusters.Add(key, cluster);
                return cluster;
            }
        }

        /// <summary>
        /// 将所有Cluster序列化为键值对形式的JSON字符串。
        /// </summary>
        /// <returns>包含所有Cluster的键值对JSON字符串。</returns>
        public Dictionary<string, string> toJSON()
        {
            Dictionary<string, string> data = new Dictionary<string, string>();

            foreach (var item in _clusters)
            {
                var cluster = item.Value;
                var json = cluster.Save();
                data.Add(item.Key, json);
            }

            return data;
        }

        /// <summary>
        /// 从JSON字符串加载指定键的Cluster数据。
        /// </summary>
        /// <param name="key">Cluster的键。</param>
        /// <param name="json">JSON字符串。</param>
        public void fromJSON(string key, string json)
        {
            if (!string.IsNullOrEmpty(json))
            {
                if (_clusters.ContainsKey(key))
                {
                    var cluster = _clusters[key];
                    cluster.Load(json);
                }
                else
                {
                    Cluster cluster = new Cluster();
                    cluster.Load(json);
                    _clusters.Add(key, cluster);
                }
            }
        }

        /// <summary>
        /// 获取指定键的Cluster实例。
        /// </summary>
        /// <param name="key">Cluster的键。</param>
        /// <returns></returns>
        public Cluster getCluster(string key) {
            if (_clusters.ContainsKey(key))
            {
                return _clusters[key];
            }
            else {
                return null;
            }
            
        }

        /// <summary>
        /// 将嵌入向量字符串转换为浮点数组。
        /// </summary>
        /// <param name="embedding">逗号分隔的浮点数</param>
        /// <returns></returns>
        public static float[] ConvertVectors(string embedding)
        {
            if (string.IsNullOrEmpty(embedding))
                return null;
         
            embedding= embedding.Trim('[', ']'); // 去除方括号   
            embedding = embedding.Replace(" ", ""); // 去除空格
            embedding = embedding.Replace("\r\n", ""); // 去除换行符
            
            var arr = embedding.Split(',');
            float[] result = new float[arr.Length];
            for (int i = 0; i < arr.Length; i++)
            {
                result[i] = float.Parse(arr[i]);
            }
            return result;
        }

        public class Cluster
        {
            private const int DEMENSIONS = 384;

            private IndexFlatL2 _index;
            private Dictionary<int, float[]> _vectorMap = new();

            public Cluster()
            {
                _index = new IndexFlatL2(DEMENSIONS);
            }

            /// <summary>
            /// 向Cluster中添加文本及其嵌入向量。
            /// </summary>
            /// <param name="id">文本的唯一标识符。</param>
            /// <param name="embedding">文本的嵌入向量。</param>
            public void AddVector(int id, float[] embedding)
            {
                _vectorMap[id] = embedding;
                RebuildIndex();
            }


            /// <summary>
            /// 从Cluster中删除文本及其嵌入向量。
            /// </summary>
            /// <param name="id">文本的唯一标识符。</param>
            public void RemoveVector(int id)
            {
                if (_vectorMap.ContainsKey(id))
                {
                    _vectorMap.Remove(id);
                    RebuildIndex();
                }
            }

            /// <summary>
            /// 重建索引以更新Cluster中的数据。
            /// </summary>
            private void RebuildIndex()
            {
                _index = new();
                if (_vectorMap.Count > 0)
                {
                    var vectors = _vectorMap.Values.ToList();
                    _index.Add(vectors);
                }
            }

            /// <summary>
            /// 搜索与给定嵌入向量最相似的文本。
            /// </summary>
            /// <param name="queryEmbedding">查询的嵌入向量。</param>
            /// <param name="topK">返回的最相似结果数量。</param>
            /// <returns>包含文本ID和相似度分数的列表。</returns>
            public List<(int id, float score)> Search(float[] queryEmbedding, int topK)
            {
                var idList = _vectorMap.Keys.ToList();
                var result = _index.Search([queryEmbedding], topK).ToArray();
                return result.Select(r =>
                {
                    int id = idList[(int)r.Label];
                    return (id, r.Distance);
                }).ToList();
            }

            /// <summary>
            /// 将Cluster数据序列化为JSON字符串。
            /// </summary>
            /// <returns>序列化后的JSON字符串。</returns>
            public string Save()
            {
                return JsonSerializer.Serialize(_vectorMap.ToDictionary(kv => kv.Key.ToString(), kv => kv.Value));
            }

            /// <summary>
            /// 从JSON字符串加载Cluster数据。
            /// </summary>
            /// <param name="json">JSON字符串。</param>
            /// <exception cref="InvalidOperationException">当JSON格式错误时抛出异常。</exception>
            public void Load(string json)
            {
                var data = JsonSerializer.Deserialize<Dictionary<string, float[]>>(json);
                if (data == null)
                    throw new InvalidOperationException("Cluster data format error, Dictionary<string, float[]> was expected.");

                _vectorMap.Clear();
                foreach (var kv in data)
                {
                    _vectorMap[int.Parse(kv.Key)] = kv.Value;
                }

                RebuildIndex();
            }
        }
    }
}
